"""IP地址分配器."""

import ipaddress
from typing import Optional

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from netsphere.models.ipam import IPAddress, IPNetwork


class IPAllocator:
    """IP地址分配器，负责IP地址的分配和回收."""
    
    def __init__(self, session: AsyncSession):
        """初始化IP分配器.
        
        Args:
            session: 数据库会话
        """
        self.session = session
    
    async def allocate_ip(
        self,
        network_id: int,
        hostname: str = None,
        description: str = None,
        mac_address: str = None,
        specific_ip: str = None,
    ) -> Optional[IPAddress]:
        """分配IP地址.
        
        Args:
            network_id: 网络ID
            hostname: 主机名
            description: 描述信息
            mac_address: MAC地址
            specific_ip: 指定IP地址
            
        Returns:
            分配的IP地址或None
        """
        network = await self.session.get(IPNetwork, network_id)
        if not network:
            raise ValueError(f"Network {network_id} not found")
        
        net = ipaddress.ip_network(f"{network.network}/{network.prefix_length}")
        
        if specific_ip:
            # 分配指定IP
            if not self._is_ip_in_network(specific_ip, net):
                raise ValueError(f"IP {specific_ip} not in network {net}")
            
            # 检查是否已存在
            existing = await self.session.execute(
                select(IPAddress).where(
                    IPAddress.address == specific_ip,
                    IPAddress.network_id == network_id
                )
            )
            if existing.scalar_one_or_none():
                raise ValueError(f"IP {specific_ip} already exists")
            
            ip = IPAddress(
                address=specific_ip,
                network_id=network_id,
                hostname=hostname,
                description=description,
                mac_address=mac_address,
                status="assigned"
            )
        else:
            # 自动分配
            available_ip = await self._find_available_ip(network_id, net)
            if not available_ip:
                return None
            
            ip = IPAddress(
                address=str(available_ip),
                network_id=network_id,
                hostname=hostname,
                description=description,
                mac_address=mac_address,
                status="assigned"
            )
        
        self.session.add(ip)
        
        # 更新网络使用统计
        network.used_addresses += 1
        self.session.add(network)
        
        await self.session.commit()
        return ip
    
    async def release_ip(self, ip_id: int) -> bool:
        """释放IP地址.
        
        Args:
            ip_id: IP地址ID
            
        Returns:
            是否成功释放
        """
        ip = await self.session.get(IPAddress, ip_id)
        if not ip:
            return False
        
        if ip.status == "assigned":
            ip.status = "available"
            ip.hostname = None
            ip.description = None
            ip.mac_address = None
            
            # 更新网络使用统计
            network = await self.session.get(IPNetwork, ip.network_id)
            if network:
                network.used_addresses = max(0, network.used_addresses - 1)
                self.session.add(network)
            
            self.session.add(ip)
            await self.session.commit()
            return True
        
        return False
    
    async def reserve_ip(
        self,
        network_id: int,
        ip_address: str,
        description: str = None,
    ) -> IPAddress:
        """保留IP地址.
        
        Args:
            network_id: 网络ID
            ip_address: IP地址
            description: 描述信息
            
        Returns:
            保留的IP地址
        """
        network = await self.session.get(IPNetwork, network_id)
        if not network:
            raise ValueError(f"Network {network_id} not found")
        
        net = ipaddress.ip_network(f"{network.network}/{network.prefix_length}")
        if not self._is_ip_in_network(ip_address, net):
            raise ValueError(f"IP {ip_address} not in network {net}")
        
        # 检查是否已存在
        existing = await self.session.execute(
            select(IPAddress).where(
                IPAddress.address == ip_address,
                IPAddress.network_id == network_id
            )
        )
        if existing_ip := existing.scalar_one_or_none():
            if existing_ip.status == "available":
                existing_ip.status = "reserved"
                existing_ip.description = description
                self.session.add(existing_ip)
                await self.session.commit()
                return existing_ip
            else:
                raise ValueError(f"IP {ip_address} is already {existing_ip.status}")
        
        ip = IPAddress(
            address=ip_address,
            network_id=network_id,
            description=description,
            status="reserved"
        )
        self.session.add(ip)
        await self.session.commit()
        return ip
    
    async def _find_available_ip(
        self,
        network_id: int,
        network: ipaddress.IPv4Network | ipaddress.IPv6Network,
    ) -> Optional[ipaddress.IPv4Address | ipaddress.IPv6Address]:
        """查找可用的IP地址.
        
        Args:
            network_id: 网络ID
            network: 网络对象
            
        Returns:
            可用的IP地址或None
        """
        # 获取已使用的地址
        result = await self.session.execute(
            select(IPAddress.address)
            .where(IPAddress.network_id == network_id)
        )
        used_ips = {row[0] for row in result.all()}
        
        # 遍历网络地址
        hosts = list(network.hosts())
        
        # 跳过网络地址和广播地址（IPv4）
        if network.version == 4:
            if len(hosts) > 0:
                hosts = hosts[1:-1] if len(hosts) > 2 else hosts
        
        # 查找第一个可用的地址
        for host in hosts:
            if str(host) not in used_ips:
                return host
        
        return None
    
    def _is_ip_in_network(
        self,
        ip_str: str,
        network: ipaddress.IPv4Network | ipaddress.IPv6Network,
    ) -> bool:
        """检查IP地址是否在网络中.
        
        Args:
            ip_str: IP地址字符串
            network: 网络对象
            
        Returns:
            是否在网络中
        """
        try:
            ip = ipaddress.ip_address(ip_str)
            return ip in network
        except ValueError:
            return False